import tensorflow.keras.models as KM
import tensorflow as tf
import warnings
from keras.applications.inception_v3 import InceptionV3
from keras import backend as K
import numpy as np
class ImageWandbLogger(tf.keras.callbacks.Callback):
def __init__(self, num_img_per_class=5, latent_dim=128):
self.num_img_per_class = num_img_per_class
self.latent_dim = latent_dim
def on_epoch_end(self, epoch, logs=None):
corelabels = tf.range(0,10 )
labels = tf.keras.utils.to_categorical(tf.repeat(corelabels, self.num_img_per_class, axis = 0 ) , 10 )
random_latent_vectors = tf.random.normal(shape=(self.num_img_per_class*10, self.latent_dim))
generated_images = self.model.generator(tf.concat([random_latent_vectors ,labels ] , axis = 1 ))
generated_images = (generated_images * 127.5) + 127.5
images = []
for i in range(self.num_img_per_class * 10 ):
img = generated_images[i].numpy()
img = tf.keras.preprocessing.image.array_to_img(img)
images.append(wandb.Image(img, caption = f'img:{i}\nepoch = {epoch}' ))
try:
wandb.log({'Generated Images' : images})
print("Successfully logged images")
except Exception as e:
print(e)
print('Logging failed')
def resize(img):
return tf.image.resize(img, (75,75,3))
class FIDCallBack(tf.keras.callbacks.Callback):
def on_epoch_start(self, epoch, logs=None):
model = self.model
try:
next(real_images)
resized = realimages.map(resize)
except:
resized = resize(real_images)
# real_images = model.images.map()
noise = tf.random.normal(shape=(128, 100))
fid_score = FID( model.generator , (-1,1 ))((real_images, noise))
self.model.fid_history['epoch'].append(epoch)
self.model.fid_history['fidscore'].append(fid_score)
if logs:
logs['fidscore'] = fid_score
print("End epoch {} of training FID is : {}".format(epoch, fid_score))
from tensorflow import keras
from tensorflow.keras import layers
class KID(keras.metrics.Metric):
#https://keras.io/examples/generative/gan_ada/
def __init__(self, name="kid", **kwargs):
super().__init__(name=name, **kwargs)
# KID is estimated per batch and is averaged across batches
self.kid_tracker = keras.metrics.Mean()
# a pretrained InceptionV3 is used without its classification layer
# transform the pixel values to the 0-255 range, then use the same
# preprocessing as during pretraining
self.encoder = keras.Sequential(
[
layers.InputLayer(input_shape=(32, 32, 3)),
# layers.Rescaling(255.0),
layers.Resizing(height=75, width=75),
layers.Lambda(keras.applications.inception_v3.preprocess_input),
keras.applications.InceptionV3(
include_top=False,
input_shape=(75, 75, 3),
weights="imagenet",
),
layers.GlobalAveragePooling2D(),
],
name="inception_encoder",
)
def polynomial_kernel(self, features_1, features_2):
feature_dimensions = tf.cast(tf.shape(features_1)[1], dtype=tf.float32)
return (features_1 @ tf.transpose(features_2) / feature_dimensions + 1.0) ** 3.0
def update_state(self, real_images, generated_images, sample_weight=None):
real_features = self.encoder(real_images, training=False)
generated_features = self.encoder(generated_images, training=False)
# compute polynomial kernels using the two sets of features
kernel_real = self.polynomial_kernel(real_features, real_features)
kernel_generated = self.polynomial_kernel(
generated_features, generated_features
)
kernel_cross = self.polynomial_kernel(real_features, generated_features)
# estimate the squared maximum mean discrepancy using the average kernel values
batch_size = tf.shape(real_features)[0]
batch_size_f = tf.cast(batch_size, dtype=tf.float32)
mean_kernel_real = tf.reduce_sum(kernel_real * (1.0 - tf.eye(batch_size))) / (
batch_size_f * (batch_size_f - 1.0)
)
mean_kernel_generated = tf.reduce_sum(
kernel_generated * (1.0 - tf.eye(batch_size))
) / (batch_size_f * (batch_size_f - 1.0))
mean_kernel_cross = tf.reduce_mean(kernel_cross)
kid = mean_kernel_real + mean_kernel_generated - 2.0 * mean_kernel_cross
# update the average KID estimate
self.kid_tracker.update_state(kid)
def result(self):
return self.kid_tracker.result()
def reset_state(self):
self.kid_tracker.reset_state()
class ImageLogger(tf.keras.callbacks.Callback):
def __init__(self, num_row=10 , num_col = 5, latent_dim=128,interval = 5):
self.num_rows = num_row
self.latent_dim = latent_dim
self.num_cols = num_col
self.interval= interval
def on_epoch_end(self, epoch, logs=None):
if epoch % self.interval == 0:
random_latent_vectors = tf.random.normal(shape=(self.num_rows * self.num_cols , self.latent_dim))
generated_images = self.model.generator([random_latent_vectors ,tf.repeat(tf.range(0,10) ,5 , axis = 0 ) ]).numpy()
generated_images = (generated_images +1 ) /2
plt.figure(figsize=(self.num_cols * 1.0, self.num_rows * 1.0))
for row in range(self.num_rows):
for col in range(self.num_cols):
index = row * self.num_cols + col
# print(generated_images[index] )
plt.subplot(self.num_rows, self.num_cols, index + 1)
plt.imshow(generated_images[index] )
plt.axis("off")
plt.tight_layout()
plt.show()
plt.close()
import os
import shutil
class WandbModelLogger(tf.keras.callbacks.Callback):
def on_epoch_end(self, epoch, logs=None):
# print('logs', logs)
if isinstance( logs, dict):
wandb.log( {**logs, **dict( epoch = epoch )})
if 'val_kid' in logs:
disstr = f"Discriminator-Epoch{epoch},KID={logs['val_kid']}"
genstr = f"Generator-Epoch{epoch},KID={logs['val_kid']}"
else:
disstr = f"Discriminator-Epoch{epoch}"
genstr = f"Generator-Epoch{epoch}"
self.model.discriminator.save(disstr,save_format="tf")
self.model.generator.save(genstr , save_format="tf")
# self.model.save(self.filepath[:-3], overwrite=True, save_format="tf")
# Log the model as artifact.
dis_name = wandb.util.make_artifact_name_safe(f"Discriminator-{wandb.run.name}")
gen_name = wandb.util.make_artifact_name_safe(f"Generator-{wandb.run.name}")
dis_artifact = wandb.Artifact(dis_name, type="model")
gen_artifact = wandb.Artifact(gen_name, type="model")
dis_artifact.add_dir(disstr)
gen_artifact.add_dir(genstr)
wandb.run.log_artifact(dis_artifact, aliases=["latest", f"Discriminator_epoch_{epoch}"])
wandb.run.log_artifact(gen_artifact, aliases=["latest", f"Generator_epoch_{epoch}"])
# Remove the SavedModel from wandb dir as we don't want to log it to save memory.
shutil.rmtree(disstr)
shutil.rmtree(genstr)
#https://medium.com/dive-into-ml-ai/customization-of-model-fit-to-write-a-combined-discriminator-generator-gan-trainer-in-keras-524bce10cf66
class GAN(KM.Model):
def __init__(self, discriminator, generator , latdim , discriminator_extra_steps =4 , img_size = 32, num_classes = 10, *args):
super().__init__()
self.discriminator = discriminator
self.generator= generator
self.latent_dim = latdim
self.discriminator_extra_steps = tf.constant(discriminator_extra_steps)
self.image_size = img_size
self.num_classes = num_classes
self.counter = None
self.gen_loss = None
self.fid_history = dict(epoch = list() , fidscore = list() )
self.d_loss_tracker = keras.metrics.Mean()
self.g_loss_tracker = keras.metrics.Mean()
def compile(self, optimizers, loss):
super().compile()
# Optimizer for discriminator
self.d_optimizer = optimizers["d_optimizer"]
# Optimizer for generator
self.g_optimizer = optimizers["g_optimizer"]
# Loss for the main objective
self.g_loss = loss["g_loss"]
# Discriminator loss
self.d_loss = loss["d_loss"]
self.kid = KID()
# self.metrics = metrics
def fit(self, *args, **kwargs):
self.img = args[0] if len(args) >= 1 else kwargs['x']
self.fid_history['epoch'] = list()
self.fid_history['fidscore'] = list()
return super().fit( *args, **kwargs)
# def forward(self, discriminator_prediction_real: torch.Tensor,
# discriminator_prediction_fake: torch.Tensor) -> torch.Tensor:
# """
# Forward pass to compute the regularization
# :param discriminator_prediction_real: (torch.Tensor) Real prediction of the discriminator
# :param discriminator_prediction_fake: (torch.Tensor) Fake prediction of the discriminator
# """
# regularization_loss = (discriminator_prediction_real - HYPERPARAMETERS["rlc_af"]).norm(dim=-1).pow(2).mean() \
# + (discriminator_prediction_fake - HYPERPARAMETERS["rlc_ar"]).norm(dim=-1).pow(2).mean()
# return HYPERPARAMETERS["rlc_w"] * regularization_loss
# def RLC_regularisation(predreal, predfake):
# param = tf.mean(tf.math.power((predreal - self.rlc_real_param ).norm(),2)) + tf.mean(tf.math.power((predfake - self.rlc_fake_param ).norm(),2)))
# return param*self.rlc_factor
def test_step(self, real_images):
if isinstance(real_images, tuple):
labels = real_images[1]
real_images = real_images[0]
# labels =
batch_size = tf.shape(real_images)[0]
latent_samples = tf.random.normal(shape=(batch_size, self.latent_dim) )
generated_images = self.generator([latent_samples , labels], training=False)
self.kid.update_state(real_images, (generated_images+1)/2)
# only KID is measured during the evaluation phase for computational efficiency
return {self.kid.name: self.kid.result()}
def train_step(self, data):
#https://keras.io/examples/generative/conditional_gan/
real_images, class_labels = data
batch_size = tf.shape(real_images)[0]
random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim) )
generated_images = self.generator([random_latent_vectors , class_labels], training = True)
with tf.GradientTape() as tape:
predictions_real = self.discriminator([real_images , class_labels], training = True)
predictions_false = self.discriminator([generated_images , class_labels] , training = True)
d_loss = self.d_loss(batch_size, predictions_real ,predictions_false , real_images, generated_images, class_labels , self.discriminator )
grads = tape.gradient(d_loss, self.discriminator.trainable_weights)
# grads = self.d_optimizer.get_unscaled_gradients( scaled_grads)
self.d_optimizer.apply_gradients(
zip(grads, self.discriminator.trainable_weights)
)
self.d_loss_tracker.update_state(d_loss)
del grads , tape
# Sample random points in the latent space
#train discriminator
if self.counter is None:
self.counter = tf.Variable(0 , trainable =False)
if self.gen_loss is None:
self.gen_loss = tf.Variable(0.0, trainable =False)
def inner_train_gen():
random_latent_vectors = tf.random.normal(shape=(batch_size, self.latent_dim))
with tf.GradientTape() as tape:
fakeimgs = self.generator([random_latent_vectors , class_labels ], training = True)
predictions = self.discriminator([fakeimgs , class_labels], training = True )
g_loss = self.g_loss( predictions)
self.gen_loss.assign(g_loss)
grads = tape.gradient(g_loss, self.generator.trainable_weights)
self.g_optimizer.apply_gradients(zip(grads, self.generator.trainable_weights))
self.g_loss_tracker.update_state(self.gen_loss)
del predictions, grads , tape
tf.cond( self.counter % self.discriminator_extra_steps ==0 , inner_train_gen , lambda :None )
self.counter.assign_add(1)
return {"d_loss": self.d_loss_tracker.result(), "g_loss": self.g_loss_tracker.result()}
from tensorflow.keras.layers import Embedding
def downsample_block(x,filters):
x = Conv2D(filters, 5, strides=2, padding='same', use_bias=True)(x)#larger kernel size of 5 avoids overfitting
x = LeakyReLU(0.2)(x)
x = Dropout(0.3)(x)#use dropout instead of bn to regularise the wgangp discriminator
return x
def discriminator():
inputimage = Input(shape=(32,32,3))
x = downsample_block(inputimage,64)
x = downsample_block(x,128)
x = downsample_block(x,256)
x = downsample_block(x,512)
label = Input(shape=(1, ))
embedded_label = Flatten()(Embedding(10, 2048)(label))
x = Flatten()(x)
x = Multiply()([x, embedded_label])
x = Dropout(0.2)(x)
x = Dense(1)(x)
return Model([inputimage,label], x)
def upsample_block(x,filters):
x = UpSampling2D(2)(x)
x = Conv2D(filters, 3, strides=1, padding='same', use_bias=False)(x)
x = BatchNormalization()(x)
x = LeakyReLU(0.2)(x)
return x
def generator():
latent_dim = 128
inputnoise = Input(shape=(latent_dim,))
label = Input(shape=(1, ))
embbededlabel = Flatten()(Embedding(10, latent_dim)(label))
model_input = Multiply()([inputnoise, embbededlabel])
x = Dense(4 * 4 * latent_dim, use_bias=False)(model_input)
x = BatchNormalization()(x)
x = LeakyReLU(0.2)(x)
x = Reshape((4, 4, latent_dim))(x)
x = upsample_block(x,128)
x = upsample_block(x, 64 )
x= UpSampling2D(2)(x)
x= Conv2D( 3,(3,3), strides = 1 , padding = 'same', use_bias = False )(x)
x= BatchNormalization()(x)
x = Activation('tanh')(x)
return Model([inputnoise,label], x)
